#!/usr/bin/env python3
"""
数据库迁移脚本：添加权限管理功能
"""

import sqlite3
import os
from datetime import datetime

def migrate_database():
    """执行数据库迁移"""
    db_path = "filemanager.db"
    
    if not os.path.exists(db_path):
        print("数据库文件不存在，请先启动应用程序创建数据库")
        return
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("开始数据库迁移...")
        
        # 1. 为files表添加新字段
        print("1. 为files表添加level和share_level字段...")
        try:
            cursor.execute("ALTER TABLE files ADD COLUMN level INTEGER")
            print("   - 添加level字段成功")
        except sqlite3.OperationalError as e:
            if "duplicate column name" in str(e):
                print("   - level字段已存在，跳过")
            else:
                raise
        
        try:
            cursor.execute("ALTER TABLE files ADD COLUMN share_level VARCHAR(20) DEFAULT 'private'")
            print("   - 添加share_level字段成功")
        except sqlite3.OperationalError as e:
            if "duplicate column name" in str(e):
                print("   - share_level字段已存在，跳过")
            else:
                raise
        
        # 2. 创建文件夹部门权限表
        print("2. 创建folder_department_permissions表...")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS folder_department_permissions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                folder_id INTEGER NOT NULL,
                department_id INTEGER NOT NULL,
                permission_type VARCHAR(20) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (folder_id) REFERENCES files (id),
                FOREIGN KEY (department_id) REFERENCES departments (id),
                UNIQUE(folder_id, department_id, permission_type)
            )
        """)
        print("   - folder_department_permissions表创建成功")
        
        # 3. 创建文件夹用户权限表
        print("3. 创建folder_user_permissions表...")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS folder_user_permissions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                folder_id INTEGER NOT NULL,
                user_id INTEGER NOT NULL,
                permission_type VARCHAR(20) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (folder_id) REFERENCES files (id),
                FOREIGN KEY (user_id) REFERENCES users (id),
                UNIQUE(folder_id, user_id, permission_type)
            )
        """)
        print("   - folder_user_permissions表创建成功")
        
        # 4. 更新现有文件的层级信息
        print("4. 更新现有文件的层级信息...")
        
        # 首先更新根目录文件夹（parent_id为NULL）
        cursor.execute("""
            UPDATE files 
            SET level = 0 
            WHERE parent_id IS NULL AND file_type = 'folder' AND level IS NULL
        """)
        
        # 递归更新子文件夹的层级
        max_iterations = 10  # 防止无限循环
        for i in range(max_iterations):
            cursor.execute("""
                UPDATE files 
                SET level = (
                    SELECT parent.level + 1 
                    FROM files parent 
                    WHERE parent.id = files.parent_id
                )
                WHERE files.parent_id IS NOT NULL 
                AND files.file_type = 'folder' 
                AND files.level IS NULL
                AND EXISTS (
                    SELECT 1 FROM files parent 
                    WHERE parent.id = files.parent_id 
                    AND parent.level IS NOT NULL
                )
            """)
            
            if cursor.rowcount == 0:
                break
        
        print(f"   - 完成{i+1}轮层级更新")
        
        # 5. 更新现有文件的共享级别
        print("5. 更新现有文件的共享级别...")
        cursor.execute("""
            UPDATE files 
            SET share_level = 'private' 
            WHERE share_level IS NULL
        """)
        print("   - 所有文件的共享级别已设置为private")
        
        # 6. 为file_shares表添加share_level字段（如果存在）
        print("6. 检查file_shares表...")
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='file_shares'")
        if cursor.fetchone():
            try:
                cursor.execute("ALTER TABLE file_shares ADD COLUMN share_level VARCHAR(20) DEFAULT 'custom'")
                print("   - 为file_shares表添加share_level字段成功")
            except sqlite3.OperationalError as e:
                if "duplicate column name" in str(e):
                    print("   - file_shares表的share_level字段已存在，跳过")
                else:
                    raise
        else:
            print("   - file_shares表不存在，跳过")
        
        # 提交所有更改
        conn.commit()
        print("\n数据库迁移完成！")
        
        # 显示迁移结果统计
        cursor.execute("SELECT COUNT(*) FROM files WHERE file_type = 'folder'")
        folder_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM files WHERE file_type = 'file'")
        file_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM folder_department_permissions")
        dept_perm_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM folder_user_permissions")
        user_perm_count = cursor.fetchone()[0]
        
        print(f"\n统计信息：")
        print(f"- 文件夹数量: {folder_count}")
        print(f"- 文件数量: {file_count}")
        print(f"- 部门权限记录: {dept_perm_count}")
        print(f"- 用户权限记录: {user_perm_count}")
        
    except Exception as e:
        print(f"迁移过程中发生错误: {e}")
        conn.rollback()
        raise
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_database()